#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#python monitor_ocr_server.py Test.png
"""
纯代码实现，无需PATH配置，适合容器化部署
"""

import os
import sys
import json
import logging
from datetime import datetime
from pathlib import Path
import cv2
import numpy as np
from PIL import Image

class MonitorTimestampOCR:

    
    def __init__(self, log_level=logging.INFO):
        """初始化"""
        self.logger = self._setup_logger(log_level)
        self.ocr_engine = None
        self.engine_type = None
        self._setup_environment()
    
    def _setup_logger(self, level):
        """设置日志"""
        logger = logging.getLogger('MonitorOCR')
        logger.setLevel(level)
        
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        
        return logger
    
    def _setup_environment(self):
        """设置环境变量，避免PATH依赖"""
        # CPU运行环境
        os.environ['FLAGS_use_gpu'] = 'false'
        os.environ['FLAGS_allocator_strategy'] = 'naive_best_fit'
        os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
        
        # 禁用不必要的警告
        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
        
        self.logger.info("环境变量配置完成")
    
    def initialize(self):
        """初始化OCR引擎"""
        try:
            from paddleocr import PaddleOCR
            
            self.logger.info("正在初始化PaddleOCR...")
            
            # 服务器优化配置
            self.ocr_engine = PaddleOCR(
                use_angle_cls=True,
                lang='ch',
                use_gpu=False,
                show_log=False,
                
                # 性能优化
                det_limit_side_len=960,
                det_limit_type='max',
                rec_batch_num=6,
                
                # 稳定性配置
                precision='fp32',
                cpu_threads=4
            )
            
            self.engine_type = 'paddleocr'
            self.logger.info("✅ PaddleOCR初始化成功")
            return True
            
        except Exception as e:
            self.logger.error(f"❌ PaddleOCR初始化失败: {e}")
            
            # 尝试EasyOCR备选
            try:
                import easyocr
                self.logger.info("尝试EasyOCR...")
                
                self.ocr_engine = easyocr.Reader(['ch_sim', 'en'], gpu=False, verbose=False)
                self.engine_type = 'easyocr'
                self.logger.info("✅ EasyOCR初始化成功")
                return True
                
            except Exception as e2:
                self.logger.error(f"❌ 所有OCR引擎初始化失败: {e2}")
                return False
    
    def preprocess_image(self, image_path):
        """监控图像预处理 - 使用直方图均衡化"""
        try:
            image = cv2.imread(image_path)
            if image is None:
                return None
            
            # 转换为灰度
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            
            # 直方图均衡化 - 改善对比度，让文字更清晰
            equalized = cv2.equalizeHist(gray)
            
            return equalized
            
        except Exception as e:
            self.logger.error(f"图像预处理失败: {e}")
            return None
    
    def extract_text(self, image_path):
        """提取文本"""
        if not self.ocr_engine:
            return {"error": "OCR引擎未初始化"}
        
        try:
            # 先尝试原始图像
            results = []
            
            if self.engine_type == 'paddleocr':
                # 首先用原始图像
                ocr_results = self.ocr_engine.ocr(image_path, cls=True)
                if ocr_results and ocr_results[0]:
                    for item in ocr_results[0]:
                        if item:
                            bbox, (text, confidence) = item[0], item[1]
                            results.append({
                                'text': text.strip(),
                                'confidence': float(confidence),
                                'bbox': bbox,
                                'source': 'original'
                            })
                
                # 如果原始图像没有结果，再尝试预处理
                if not results:
                    processed = self.preprocess_image(image_path)
                    if processed is not None:
                        temp_path = "temp_processed.jpg"
                        cv2.imwrite(temp_path, processed)
                        
                        ocr_results = self.ocr_engine.ocr(temp_path, cls=True)
                        if ocr_results and ocr_results[0]:
                            for item in ocr_results[0]:
                                if item:
                                    bbox, (text, confidence) = item[0], item[1]
                                    results.append({
                                        'text': text.strip(),
                                        'confidence': float(confidence),
                                        'bbox': bbox,
                                        'source': 'processed'
                                    })
                        
                        # 清理临时文件
                        if os.path.exists(temp_path):
                            os.remove(temp_path)
            
            elif self.engine_type == 'easyocr':
                ocr_results = self.ocr_engine.readtext(image_path)
                for bbox, text, confidence in ocr_results:
                    results.append({
                        'text': text.strip(),
                        'confidence': float(confidence),
                        'bbox': bbox,
                        'source': 'original'
                    })
            
            return {"success": True, "results": results}
            
        except Exception as e:
            self.logger.error(f"文本提取失败: {e}")
            return {"error": str(e)}
    
    def calculate_timestamp_score(self, text):
        """计算时间戳相关性得分 - 不依赖固定格式"""
        score = 0
        
        # 数字密度评分（更灵活）
        digits = sum(1 for c in text if c.isdigit())
        if digits > 0:
            digit_ratio = digits / len(text)
            # 高数字密度更可能是时间戳
            if digit_ratio >= 0.8:
                score += 0.6
            elif digit_ratio >= 0.5:
                score += 0.4
            elif digit_ratio >= 0.3:
                score += 0.2
        
        # 包含可能的时间分隔符（但不强制特定格式）
        separators = [':', '-', '/', '.', ' ', '年', '月', '日', '时', '分', '秒']
        sep_count = sum(1 for sep in separators if sep in text)
        score += min(sep_count * 0.1, 0.3)
        
        # 长度评分（监控时间戳长度通常在这个范围）
        text_len = len(text)
        if 8 <= text_len <= 25:  # 完整时间戳长度
            score += 0.3
        elif 4 <= text_len <= 7:   # 可能的部分时间戳
            score += 0.2
        elif 3 <= text_len <= 3:   # 短数字片段
            score += 0.15
        
        # 数字连续性检查（时间戳通常有连续数字）
        import re
        consecutive_digits = re.findall(r'\d+', text)
        if consecutive_digits:
            max_consecutive = max(len(num) for num in consecutive_digits)
            if max_consecutive >= 4:
                score += 0.2
            elif max_consecutive >= 2:
                score += 0.1
        
        # 常见时间数字范围检查（不限制格式）
        if text.isdigit():
            num = int(text)
            if 2000 <= num <= 2099:    # 年份
                score += 0.4
            elif 1 <= num <= 31:       # 日期
                score += 0.2
            elif 0 <= num <= 23:       # 小时
                score += 0.2
            elif 0 <= num <= 59:       # 分钟/秒
                score += 0.2
        
        return min(score, 1.0)
    
    def recognize_timestamp(self, image_path):
        """识别监控时间戳 - 主接口"""
        self.logger.info(f"开始识别: {os.path.basename(image_path)}")
        
        if not os.path.exists(image_path):
            return {"error": f"图像文件不存在: {image_path}"}
        
        # 提取文本
        result = self.extract_text(image_path)
        if "error" in result:
            return result
        
        # 过滤和评分
        candidates = []
        for item in result['results']:
            text = item['text']
            confidence = item['confidence']
            
            # 基本长度过滤（允许更短的文本，因为可能是部分时间戳）
            if len(text.strip()) < 2:
                continue
            
            timestamp_score = self.calculate_timestamp_score(text)
            total_score = confidence * 0.6 + timestamp_score * 0.4
            
            # 降低阈值，允许更多候选（包括部分时间戳）
            if total_score > 0.2:
                candidates.append({
                    'text': text,
                    'confidence': confidence,
                    'bbox': item['bbox'],
                    'timestamp_score': timestamp_score,
                    'total_score': total_score
                })
        
        # 按得分排序
        candidates.sort(key=lambda x: x['total_score'], reverse=True)
        
        return {
            "success": True,
            "image_path": image_path,
            "engine": self.engine_type,
            "total_detected": len(candidates),
            "candidates": candidates,
            "timestamp": datetime.now().isoformat()
        }
    
    def batch_process(self, image_paths):
        """批量处理"""
        results = {}
        for image_path in image_paths:
            results[image_path] = self.recognize_timestamp(image_path)
        return results
    
    def save_result(self, result, output_file=None):
        """保存结果"""
        if not output_file:
            output_file = f"timestamp_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(result, f, ensure_ascii=False, indent=2)
            self.logger.info(f"结果已保存: {output_file}")
            return output_file
        except Exception as e:
            self.logger.error(f"保存失败: {e}")
            return None

def main():
    """主程序"""
    print("=" * 60)
    print("监控时间戳识别系统")
    print("=" * 60)
    
    # 初始化OCR
    ocr = MonitorTimestampOCR()
    
    if not ocr.initialize():
        print("❌ OCR引擎初始化失败")
        return
    
    print(f"✅ 使用引擎: {ocr.engine_type}")
    
    # 获取图像路径
    if len(sys.argv) > 1:
        image_path = sys.argv[1]
    else:
        image_path = "Test.png"
    
    if not os.path.exists(image_path):
        print(f"❌ 图像不存在: {image_path}")
        return
    
    print(f"\n📷 正在分析: {os.path.basename(image_path)}")
    
    # 识别时间戳
    result = ocr.recognize_timestamp(image_path)
    
    if "error" in result:
        print(f"❌ 识别失败: {result['error']}")
        return
    
    # 显示结果
    print(f"\n📊 检测结果:")
    print(f"   引擎: {result['engine']}")
    print(f"   候选数量: {result['total_detected']}")
    
    if result['candidates']:
        print(f"\n🎯 时间戳候选（前3个）:")
        print("-" * 50)
        
        for i, candidate in enumerate(result['candidates'][:3], 1):
            print(f"{i}. 📅 {candidate['text']}")
            print(f"   OCR置信度: {candidate['confidence']:.3f}")
            print(f"   时间戳相关性: {candidate['timestamp_score']:.3f}")
            print(f"   综合得分: {candidate['total_score']:.3f}")
            print()
        
        # 最佳匹配
        best = result['candidates'][0]
        print(f"🏆 最佳匹配: {best['text']} (得分: {best['total_score']:.3f})")
        
        # 保存结果
        output_file = ocr.save_result(result)
        if output_file:
            print(f"💾 结果已保存: {output_file}")
        
    else:
        print("⚠️  未检测到时间戳候选")
    
    print("\n✅ 识别完成")

if __name__ == "__main__":
    main()
